import numpy as np
import torch
import matplotlib.pyplot as plt
import torch.nn as nn
from torchvision.transforms import ToTensor



def test1():
    img = plt.imread('./img.jpg')
    img2 = torch.tensor( img,dtype=torch.float32).permute(2,0,1)
    img3 = img2.unsqueeze(0)
    con = nn.Conv2d(in_channels=3,out_channels=3,kernel_size=3,stride=2,padding=0)
    output = con(img3)
    img4 = output[0].permute(1,2,0)
    feature1 = img4[:,:,0].detach().numpy()
    print('feature1', feature1.shape)
    plt.imshow(feature1)
    plt.show()

if __name__ == '__main__':
    test1()
